#define _GNU_SOURCE
#include "init.h"
#include "basics/error_handle.h"
#include "basics/macros.h"
#include "cgroup/cg.h"
#include "env/env.h"
#include "netns.h"
#include "persistence/persistence.h"
#include "security/capability.h"
#include <sched.h>
#include <string.h>
#include <sys/mount.h>
#include <sys/prctl.h>
#include <sys/stat.h>
#include <sys/syscall.h>
#include <sys/wait.h>
#include <time.h>

#define STACK_SIZE (1024 * 1024) // 1MB stack for child

static char child_stack[STACK_SIZE];

struct child_args {
    const char *root_path;
    const char *file;
    const int container_id;
    char **argv;
    char **envp;
    int c_wait_p_pipe[2];
    int p_wait_c_pipe[2];
    int netns_ip;
};

#define MAP_BUF_SIZE 128

static int child_func(void *arg);
/// this is a must have option before writing into gid_map
/// for more info, see `man user_namespaces`
static void setgroups_write(pid_t child_pid);
static int update_map(char *map_file, char *map_buf);

/// hint: cntr is the abbr of container
int cntr_create(const char *image_path, const char *command, char *argv[],
                int pipe_fd_child, int netns_ip) {
    struct metadata meta;
    srand(time(NULL));
    int cid = rand();
    meta.conatiner_id = cid;
    check(copy_img(image_path, meta.conatiner_id) == 0);
    struct child_args args = {get_root_path(cid), command, cid, &argv[0]};
    struct metadata meta_data = {cid, 0, 0, 0, 0};

    args.netns_ip = netns_ip;
    /// get env from a predefined text file
    int env_fd = open(get_env_path(cid), O_RDONLY);
    check(env_fd != -1);
    struct env env = {"\0", {NULL}};
    check(read(env_fd, env.env_buf, ENV_BUF_SIZE) != -1);
    if (env_buf_to_envp(&env, cid) == false) {
        err(EXIT_FAILURE, "[neodocker]: env is NULL");
    }
    check(close(env_fd) == 0);
    args.envp = env.envp;
    check(pipe(args.c_wait_p_pipe) != -1);
    check(pipe(args.p_wait_c_pipe) != -1);
    /// create child process to start the container
    int flags = CLONE_NEWNET | CLONE_NEWIPC | CLONE_NEWNS | CLONE_NEWCGROUP |
                CLONE_NEWPID | CLONE_NEWUSER | CLONE_NEWUTS | SIGCHLD;
    pid_t child_pid = clone(child_func, child_stack + STACK_SIZE, flags, &args);
    check(write(pipe_fd_child, &child_pid, sizeof(pid_t)) != -1);
    check(close(pipe_fd_child) != -1);
    check(child_pid >= 0);
    /// small trick to avoid one syscall(getpid)
    meta_data.root_pid = child_pid;
    char uid_map[PATH_MAX];
    char gid_map[PATH_MAX];
    snprintf(uid_map, PATH_MAX, "/proc/%d/uid_map", child_pid);
    snprintf(gid_map, PATH_MAX, "/proc/%d/gid_map", child_pid);
    char map_buf[MAP_BUF_SIZE];
    uid_t uid = getuid();
    gid_t gid = getgid();
    /// if it is root user, map all users
    if (uid == 0) {
        snprintf(map_buf, MAP_BUF_SIZE, "0 %u %u", uid, UINT32_MAX);
        snprintf(map_buf, MAP_BUF_SIZE, "0 %u %u", gid, UINT32_MAX);
    } else { /// if it is a normal user, map only itself to root
        snprintf(map_buf, MAP_BUF_SIZE, "0 %u 1", uid);
        snprintf(map_buf, MAP_BUF_SIZE, "0 %u 1", gid);
    }
    check(close(args.p_wait_c_pipe[1]) == 0);
    char ch;
    check(read(args.p_wait_c_pipe[0], &ch, sizeof(char)) == 0);
    check(close(args.p_wait_c_pipe[0]) == 0);
    check(update_map(uid_map, map_buf) == 0);
    setgroups_write(child_pid);
    check(update_map(gid_map, map_buf) == 0);
    check(store_metadata(meta_data) == 0);
    check(close(args.c_wait_p_pipe[1]) == 0);
    check(close(args.c_wait_p_pipe[0]) == 0);
    // check(waitpid(child_pid, NULL, 0) == child_pid);
    return cid;
}

int cntr_attach(const int container_id, const char *command, char *argv[]) {
    struct metadata meta_data;
    meta_data.conatiner_id = container_id;
    check(load_metadata(&meta_data) == 0);
    int env_fd = open(get_env_path(container_id), O_RDONLY);
    check(env_fd != -1);
    struct env env = {"\0", {NULL}};
    check(read(env_fd, env.env_buf, ENV_BUF_SIZE) != -1);
    if (env_buf_to_envp(&env, container_id) == false) {
        err(EXIT_FAILURE, "[neodocker]: env is NULL");
    }
    check(close(env_fd) != -1);
    /// warning: pid and pid_for_children must be set in the parent
    // const char *namespaces[] = {"pid", "pid_for_children"};
    // char ns_path[PATH_MAX];
    // for (int i = 0; i < sizeof(namespaces) / sizeof(char *); i++) {
    //   snprintf(ns_path, PATH_MAX, "/proc/%d/ns/%s", meta_data.root_pid,
    //            namespaces[i]);
    //   int ns_fd = open(ns_path, O_RDONLY);
    //   check(ns_fd != -1);
    //   check(setns(ns_fd, 0) != -1);
    //   check(close(ns_fd) != -1);
    // }
    int fd = syscall(SYS_pidfd_open, meta_data.root_pid, 0);
    check(fd != -1);
    check(setns(fd, CLONE_NEWPID) != -1);
    pid_t pid = fork();
    apply_to_cgroup(container_id, pid);
    if (pid == 0) {
        check(setns(fd, CLONE_NEWUTS | CLONE_NEWCGROUP | CLONE_NEWNS |
                            CLONE_NEWIPC | CLONE_NEWNET) != -1);
        // const char *namespaces[] = {"cgroup", "mnt", "ipc", "net", "uts"};
        // for (int i = 0; i < sizeof(namespaces) / sizeof(char *); i++) {
        //   snprintf(ns_path, PATH_MAX, "/proc/%d/ns/%s", 1, namespaces[i]);
        //   int ns_fd = open(ns_path, O_RDONLY);
        //   check(ns_fd != -1);
        //   check(setns(ns_fd, 0) != -1);
        //   check(close(ns_fd) != -1);
        // }
        // char hostname[11];
        // snprintf(hostname, sizeof(hostname), "%d", container_id);
        // check(sethostname(hostname, sizeof(hostname)) == 0);
        check(chroot(get_root_path(container_id)) == 0);
        /// chroot will not cd into the correct directory
        check(chdir("/") == 0);
        check(mount("proc", "/proc", "proc", MS_NOEXEC | MS_NOSUID | MS_NODEV,
                    NULL) == 0);
        // snprintf(ns_path, PATH_MAX, "/proc/%d/ns/%s", 1, "user");
        // int ns_fd = open(ns_path, O_RDONLY);
        // check(ns_fd != -1);
        check(setns(fd, CLONE_NEWUSER) != -1);
        // check(setns(ns_fd, 0) != -1);
        // check(close(ns_fd) != -1);
        check(setgid(0) != -1);
        check(setuid(0) != -1);
        check(execvpe(command, argv, env.envp) != -1);
        err(EXIT_FAILURE, "[neodocker]: execvpe failed");
    } else {
        check(waitpid(pid, NULL, 0) == pid);
    }
    return 0;
}

int cntr_delete(const int container_id) {
    struct metadata meta_data;
    meta_data.conatiner_id = container_id;
    check(load_metadata(&meta_data) == 0);
    /// root_pid is the group's pgid
    check(killpg(meta_data.root_pid, SIGTERM) == 0);
    check(delete_container(container_id) == 0);
    return 0;
}

static int child_func(void *arg) {
    /// https://stackoverflow.com/questions/47296408/cannot-open-uid-map-for-writing-from-an-app-with-cap-setuid-capability-set
    /// this is critical: see this url above for more information, man proc(5)
    /// and
    ///    ↓    man prctl(2) has more information
    check(prctl(PR_SET_DUMPABLE, 1) != -1);
    struct child_args *args = arg;
    check(close(args->p_wait_c_pipe[1]) == 0);
    check(close(args->p_wait_c_pipe[0]) == 0);
    check(close(args->c_wait_p_pipe[1]) == 0);
    char ch;
    if (read(args->c_wait_p_pipe[0], &ch, sizeof(char)) != 0) {
        fprintf(stderr,
                "[neodocker]: Failed in child: read from pipe returned != 0\n");
        exit(EXIT_FAILURE);
    }
    check(close(args->c_wait_p_pipe[0]) == 0);
    // check(set_full_cap() == 0);
    char hostname[11];
    snprintf(hostname, sizeof(hostname), "%d", args->container_id);
    check(sethostname(hostname, sizeof(hostname)) == 0);
    check(chroot(args->root_path) == 0);
    /// chroot will not cd into the correct directory
    check(chdir("/") == 0);
    // check(mkdir("/proc", S_ISUID | S_ISGID | S_IRUSR | S_IXUSR | S_IRGRP |
    //                          S_IXGRP | S_IROTH | S_IXOTH) == 0);
    check(mount("proc", "/proc", "proc", MS_NOEXEC | MS_NOSUID | MS_NODEV,
                NULL) == 0);
    setpgrp();
    create_netns_container(args->netns_ip);
    const char *command = "/usr/sbin/ip route add default via 10.0.0.1";
    check(system(command) == 0);
    /// capability setting for security
    check(default_cap() == 0);
    /// execvpe should not return
    check(execvpe(args->file, args->argv, args->envp) != -1);
    /// so it should not reach here
    err(EXIT_FAILURE, "[neodocker]: execvpe failed");
}

static int update_map(char *map_file, char *map_buf) {
    int fd = open(map_file, O_RDWR);
    if_fail(fd != -1, return -1);
    if_fail(write(fd, map_buf, MAP_BUF_SIZE) == MAP_BUF_SIZE, goto error);
    if_fail(close(fd) == 0, return -1);
    return 0;
error:
    check(close(fd) == 0);
    return -1;
}

static void setgroups_write(pid_t child_pid) {
    char setgroups_path[PATH_MAX];
    int fd;
    snprintf(setgroups_path, PATH_MAX, "/proc/%d/setgroups", child_pid);
    fd = open(setgroups_path, O_RDWR);
    check(fd != -1);
    check(write(fd, "deny", sizeof("deny")) != -1);
    check(close(fd) == 0);
}

#undef MAP_BUF_SIZE
